#!/usr/bin/python
# -*-coding:utf-8-*-
import os
import pandas as pd
import numpy as np

from zg_data_process.zg_data_concat import DataConcat
from zg_data_process.zg_data_process import DataProcess

from dir_info import raw_X_data_dir
from dir_info import processed_X_data_dir

data_concat_api = DataConcat(label_data_filename='processed_new_stock_label_ci1_data',
                             size_db='validation')

data_process_api = DataProcess()

def main(raw_data_filename,
         processed_data_filename):
    print(raw_data_filename, 'process...')

    # TODO - 读取数据
    raw_data = pd.read_hdf(os.path.join(raw_X_data_dir, raw_data_filename + '.h5'))

    data_columns = raw_data.columns

    raw_data = raw_data.reset_index()

    # TODO - label data
    label_user_behavior_data = data_concat_api.label_factor_data(df=raw_data,
                                                                 start_date=None,
                                                                 end_date=None,
                                                                 cache=False,
                                                                 cache_filename=None,
                                                                 add_size_factor=False,
                                                                 refresh=True,
                                                                 verbose=True,
                                                                 copy=True)

    print('process starts...')

    # TODO - 去极值
    label_user_behavior_data = data_process_api.cs_mad_outlier_process(df=label_user_behavior_data,
                                                                       columns=data_columns)
    print('mad outlier process done!')

    # # TODO - 横截面标准化
    # label_user_behavior_data = data_process_api.cs_z_score_normalization_process(df=label_user_behavior_data,
    #                                                                              columns=data_columns)
    # print('cs zscore done!')

    # TODO - 去停牌
    label_user_behavior_data = data_process_api.paused_stock_filtration(df=label_user_behavior_data)
    print('filter paused stock done!')

    # TODO - 去st
    label_user_behavior_data = data_process_api.st_filteration(df=label_user_behavior_data)
    print('filter st stock done!')

    # TODO - 去一字板
    label_user_behavior_data = data_process_api.oneline_stock_filtration(df=label_user_behavior_data)
    print('filter online stock done!')

    label_user_behavior_data = label_user_behavior_data.set_index(['date', 'stock_code'])

    label_user_behavior_data = label_user_behavior_data[data_columns]

    # label_user_behavior_data = pd.read_hdf(os.path.join(processed_X_data_dir, processed_data_filename + '.h5'))
    label_user_behavior_data.to_hdf(os.path.join(processed_X_data_dir, processed_data_filename + '.h5'),
                                    format='table',
                                    key=processed_data_filename)

    print(raw_data_filename, 'done!')


if __name__ == '__main__':
    raw_data_filename = 'X_hf_order_quantile_data'
    processed_data_filename = 'processed_X_hf_order_quantile_data'

    main(raw_data_filename=raw_data_filename,
         processed_data_filename=processed_data_filename)

